from psycopg2 import pool
from tqdm import tqdm
from datetime import datetime


class FeatureGetAndUpdate():
    """ Feature get or update on time.

    """

    def __init__(self, table='smc_sign_feature'):
        """ Parameter initialization.

        Args:
            table: string, the name of the table that saves the feature
        """
        self.table = table
        
        try:
            self.connectPool = pool.SimpleConnectionPool(2, 10, host="124.70.195.145", port="16800",
                                                         user="postgres", password="Admin@2021",
                                                         database="smc", keepalives=1,
                                                         keepalives_idle=30, keepalives_interval=10,
                                                         keepalives_count=5)
        except Exception as e:
            print(e)

    def psql_connect(self):
        """ PostgreSQL server connect.

        Returns:
            conn: get connection
        """
        conn = self.connectPool.getconn()
        cursor = conn.cursor()
        print("Connect PostgreSQL server seccuss!")
        
        return conn, cursor
    
    def closeConnect(self, conn, cursor):
        cursor.close()
        self.connectPool.putconn(conn)

    def create_table(self, delete_exist_table=False):
        """ Create a table to save features.

        Args:
            cursor:             the cursor obtained after connecting to the PostgreSQL server
            delete_exist_table: bool, delete the existing table with the same name if ture
        """
        conn, cursor = self.psql_connect()

        if delete_exist_table:
            # (慎重选择)如果存在对应的特征表格，就先删除
            cursor.execute(r"DROP TABLE IF EXISTS %s;" % self.table)

        # 创建数据表SQL语句
        sql_create_table = """
        CREATE TABLE %s (
        heart_id serial PRIMARY KEY,
        person_id varchar(50) NOT NULL,
        create_time timestamp(6),
        max_value varchar(8),
        min_value varchar(8),
        diff_value varchar(16),
        mean_value varchar(16),
        std_value varchar(16),
        update_time timestamp(6),
        total_heart_num varchar(50)
        );
        COMMENT ON COLUMN %s.heart_id IS '主键';
        COMMENT ON COLUMN %s.person_id IS '人员ID';
        COMMENT ON COLUMN %s.create_time IS '创建时间';
        COMMENT ON COLUMN %s.max_value IS '特征：最大值';
        COMMENT ON COLUMN %s.min_value IS '特征：最小值';
        COMMENT ON COLUMN %s.diff_value IS '特征：极差';
        COMMENT ON COLUMN %s.mean_value IS '特征：均值';
        COMMENT ON COLUMN %s.std_value IS '特征：标准差';
        COMMENT ON COLUMN %s.update_time IS '特征更新时间';
        COMMENT ON COLUMN %s.total_heart_num IS '心率数据总数';
        """ % (self.table, self.table, self.table, self.table, self.table, self.table, self.table, self.table, self.table, self.table, self.table)

        cursor.execute(sql_create_table)

        self.closeConnect(conn, cursor)

        print("Create new table:%s success!" % self.table)

    def get_person_id(self, cursor):
        """ Get the person id from the feature table.

        Args:
            cursor: the cursor obtained after connecting to the PostgreSQL server.

        Returns:
            all_distinct_person_id: list
        """
        cursor.execute(r"select distinct person_id from %s;" % self.table)
        all_distinct_person_id = cursor.fetchall()
        all_distinct_person_id = [i[0] for i in all_distinct_person_id]
        return all_distinct_person_id

    def update_features(self, features, person_id, update_time, total_heart_num):
        """ Update feature.

        Args:
            features:         list
            person_id:        string, person ID of feature
            create_time:      datatime, the time of creating feature table
            total_heart_num:  string, total amount of heart rate data
        """
        conn, cursor = self.psql_connect()

        cursor.execute(
            "select heart_id from %s where person_id='%s';" % (self.table, person_id))
        start_heart_id = cursor.fetchone()[0]

        print("Found features in", self.table, ", now update the feature...")
        for f in tqdm(features):
            sql_update = """
            update %s
            set max_value='%s', min_value='%s', diff_value='%s', mean_value='%s', std_value='%s', update_time='%s', total_heart_num='%s'
            where heart_id='%s'
            """ % (self.table, f[0], f[1], f[2], f[3], f[4], update_time, total_heart_num, start_heart_id)

            start_heart_id += 1
            cursor.execute(sql_update)
        print("Update features success!")

        self.closeConnect(conn, cursor)

    def get_feature(self, person_id):
        """ Get feature of the person ID.

        Args:
            person_id: string, person ID of feature

        Returns:
            features:        list 
            update_time:     datatime, the time of updating feature table
            total_heart_num: string, total amount of heart rate data
        """
        conn, cursor = self.psql_connect()

        all_distinct_person_id = self.get_person_id(cursor)

        if person_id not in all_distinct_person_id:  # create feature
            print("Not found features in", self.table, ", now initialize the feature...")

            # initialize the features with 0
            features = [[0 for _ in range(5)] for i in range(72)]
            update_time = datetime.now()
            create_time = update_time
            total_heart_num = 0

            for f in tqdm(features):
                sql_create = """
                INSERT INTO %s (person_id, create_time, max_value, min_value, diff_value, mean_value, std_value, update_time, total_heart_num)
                VALUES ('%s', '%s', '%s', '%s', '%s', '%s', '%s', '%s', '%s');
                """ % (self.table, person_id, create_time, f[0], f[1], f[2], f[3], f[4], update_time, total_heart_num)

                cursor.execute(sql_create)
            print("Initialize features success!")

        else:
            print("Found features in", self.table, ", now get the feature...")

            cursor.execute("select max_value, min_value, diff_value, mean_value, std_value from %s where person_id='%s';" %
                           (self.table, person_id))
            features = cursor.fetchall()

            cursor.execute("select update_time, total_heart_num from %s where person_id='%s';" %
                           (self.table, person_id))
            data = cursor.fetchone()
            update_time = data[-2]
            total_heart_num = data[-1]

            print("Get features success!")

        self.closeConnect(conn, cursor)

        return features, update_time, total_heart_num


if __name__ == "__main__":
    # FeatureGetAndUpdate().psql_connect()
    
    
    # Create new table if need
    FeatureGetAndUpdate().create_table(delete_exist_table=True)

    # get features
    feature, update_time, total_heart_num = FeatureGetAndUpdate().get_feature('640')
    # print(feature, update_time, total_heart_num)

    # feature demo
    features = [[i+1 for i in range(5)] for _ in range(len(feature))]
    person_id = "640"
    create_time = '2021-04-11 11:48:10.852'
    update_time = '2021-04-11 11:48:10.852'
    total_heart_num = '153'

    # 特征保存与更新
    FeatureGetAndUpdate().update_features(features, person_id,
                                          create_time, total_heart_num)
